package com.whosly.demo.jmh.rocksdb;

import org.apache.commons.io.FileUtils;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.DBOptions;
import org.rocksdb.Options;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;

/**
 * @author fengyang
 * @date 2023/8/8 下午8:37
 * @see <a href="yueny09@163.com">mailTo: yueny09@163.com</a>
 */
public class RocksDBManager
{
    private Logger logger = LoggerFactory.getLogger(this.getClass());

    public static boolean ROCKSDB_USE_DIO_READ = false;
    public static boolean ROCKSDB_USE_DIO_WRITE = false;

    private final static RocksDBManager singleInstance = new RocksDBManager();

    public static RocksDBManager getInstance() {
        return singleInstance;
    }

    protected DBOptions defaultDbOptions;
    protected ColumnFamilyOptions defaultCfOptions;

    protected RocksDBManager() {
        setOptions();
    }

    /**
     * Create and open a database instance using the predefined options. If the
     * database already exists, just open it. @see setCreateIfMissing(true).
     *
     * @param dbDir
     * @return
     */
    public RocksDBWrapper create(File dbDir) throws Exception
    {
        cleanLogDir(dbDir);

        RocksDB db = null;
        try {
            Options dbOptions = new Options(defaultDbOptions, defaultCfOptions);
            dbOptions.setCreateIfMissing(true);
            dbOptions.setDbLogDir(createLogDir(dbDir.getAbsolutePath()));

            db = RocksDB.open(dbOptions, dbDir.getAbsolutePath());
        } catch (RocksDBException e) {
            String s = "Create RocksDB database \"" + dbDir + "\" failed";
            printRocksDBException(s, e);
            throw e;
        }
        logger.info("Create RocksDB database {} OK.", dbDir);

        return new RocksDBWrapper(dbDir.getAbsolutePath(), db);
    }

    /**
     * 清空 dataDir 文件夹
     */
    private void cleanLogDir(File dataDir) throws IOException
    {
        File logDir = dataDir;

        if (logDir.exists()) {
            FileUtils.deleteDirectory(logDir);
        }
    }

    /**
     * Open a database instance using the predefined options.
     *
     * @param dbDir
     * @return
     */
    public RocksDBWrapper open(File dbDir) throws RocksDBException{
        RocksDB db = null;
        try {
            Options dbOptions = new Options(defaultDbOptions, defaultCfOptions);
            dbOptions.setCreateIfMissing(false);
            dbOptions.setDbLogDir(createLogDir(dbDir.getAbsolutePath()));

            db = RocksDB.open(dbOptions, dbDir.getAbsolutePath());
        } catch (RocksDBException e) {
            String s = "Open(readwrite) RocksDB database \"" + dbDir + "\" failed";
            printRocksDBException(s, e);
            throw e;
        }

        return new RocksDBWrapper(dbDir.getAbsolutePath(), db);
    }

    public void close(RocksDB rocksdb) {
        if (rocksdb != null) {
            rocksdb.close();
            rocksdb = null;
        }
    }

    private void setOptions() {
        ////////////////////////////////////////////////////////////////
        // defaultDbOptions
        ////////////////////////////////////////////////////////////////
        DBOptions dbOptions = new DBOptions();
        dbOptions.setCreateIfMissing(true);
        this.defaultDbOptions = dbOptions;

        ////////////////////////////////////////////////////////////////
        // defaultCfOptions
        ////////////////////////////////////////////////////////////////
        ColumnFamilyOptions cfo = new ColumnFamilyOptions();
        this.defaultCfOptions = cfo;
    }

    private void printRocksDBException(String prefix, RocksDBException e) {
        logger.error("{}, status=>{}, code=>{}, subcode=>{}", prefix, e.getStatus().getState(),
                e.getStatus().getCode().toString(), e.getStatus().getSubCode().toString());
    }

    private String createLogDir(String dbDir) {
        String logDir = dbDir;
        File fDir = new File(logDir);

        if(!fDir.exists()) {
            fDir.mkdirs();
        }

        return logDir;
    }
}
